package query

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

// CambriconQuery Cambricon MLU查询实现
type CambriconQuery struct{}

func init() {
	RegisterQueryer("cambricon", func() GPUQuery { return &CambriconQuery{} })
}

func (q *CambriconQuery) Query() ([]GPUInfo, error) {
	// 获取设备列表
	devices, err := q.getDevices()
	if err != nil {
		return nil, fmt.Errorf("获取设备列表失败: %v", err)
	}

	// 如果设备数量为0，直接返回空结果
	if len(devices) == 0 {
		return []GPUInfo{}, nil
	}

	// 构建设备ID字符串（逗号分隔）
	deviceIDs := make([]string, len(devices))
	for i, device := range devices {
		deviceIDs[i] = strconv.Itoa(device)
	}
	deviceIDStr := strings.Join(deviceIDs, ",")

	// 获取所有设备信息
	output, err := runCommand("cnmon", "info", "-c", deviceIDStr, "-u", "-m")
	if err != nil {
		return nil, fmt.Errorf("获取设备信息失败: %v", err)
	}

	// 解析所有设备信息
	gpus, err := q.parseAllDevicesInfo(output, devices)
	if err != nil {
		return nil, fmt.Errorf("解析设备信息失败: %v", err)
	}

	return gpus, nil
}

// getDevices 获取设备列表
func (q *CambriconQuery) getDevices() ([]int, error) {
	output, err := runCommand("cnmon", "-l")
	if err != nil {
		return nil, err
	}

	// 使用正则表达式快速提取设备ID
	deviceRegex := regexp.MustCompile(`MLU\s+(\d+):`)
	lines := strings.Split(output, "\n")
	var devices []int

	for _, line := range lines {
		if matches := deviceRegex.FindStringSubmatch(line); matches != nil {
			if deviceID, err := strconv.Atoi(matches[1]); err == nil {
				devices = append(devices, deviceID)
			}
		}
	}

	if len(devices) == 0 {
		return nil, fmt.Errorf("未找到有效的MLU设备")
	}

	return devices, nil
}

// parseAllDevicesInfo 解析所有设备信息
func (q *CambriconQuery) parseAllDevicesInfo(output string, devices []int) ([]GPUInfo, error) {
	// 预编译正则表达式
	deviceHeaderRegex := regexp.MustCompile(`Card\s+(\d+)`)
	utilRegex := regexp.MustCompile(`MLU Average\s+:\s+(\d+) %`)
	memUsedRegex := regexp.MustCompile(`Used\s+:\s+(\d+) MiB`)
	memTotalRegex := regexp.MustCompile(`Total\s+:\s+(\d+) MiB`)

	// 分割输出为行
	lines := strings.Split(output, "\n")
	gpus := make([]GPUInfo, len(devices))
	currentDevice := -1

	for _, line := range lines {
		// 检查是否为新设备部分
		if matches := deviceHeaderRegex.FindStringSubmatch(line); matches != nil {
			if deviceID, err := strconv.Atoi(matches[1]); err == nil {
				currentDevice = deviceID
			}
			continue
		}

		// 如果不在设备部分，继续搜索
		if currentDevice == -1 {
			continue
		}

		// 查找设备在列表中的索引
		deviceIndex := -1
		for i, d := range devices {
			if d == currentDevice {
				deviceIndex = i
				break
			}
		}

		// 如果设备不在列表中，继续搜索
		if deviceIndex == -1 {
			continue
		}

		// 解析使用率
		if matches := utilRegex.FindStringSubmatch(line); matches != nil {
			if utilization, err := strconv.Atoi(matches[1]); err == nil {
				gpus[deviceIndex].Utilization = utilization
			}
		}

		// 解析内存使用量
		if matches := memUsedRegex.FindStringSubmatch(line); matches != nil {
			if used, err := strconv.Atoi(matches[1]); err == nil {
				gpus[deviceIndex].MemoryUsed = used
			}
		}

		// 解析内存总量
		if matches := memTotalRegex.FindStringSubmatch(line); matches != nil {
			if total, err := strconv.Atoi(matches[1]); err == nil {
				gpus[deviceIndex].MemoryTotal = total
			}
		}
	}

	// 设置设备索引
	for i, device := range devices {
		gpus[i].Index = device
	}

	return gpus, nil
}
